# -*- coding: utf-8 -*-
# @Time    : 2023/5/8 16:47
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: __init__.py

import paddle
from .Unet import UNet
from .SwinTransformer import SwinTransformer_tiny_patch4_window7_224, SwinTransformer_base_patch4_window7_224
from .SwinTransformer_Target import Target_SwinTransformer_tiny_patch4_window7_224, Target_SwinTransformer_base_patch4_window7_224


def diffusion(config):
    model = eval(config["type"])(config)
    if "pretrained" in config and config["pretrained"] is not None:
        model.set_state_dict(paddle.load(config["pretrained"]))
    return model


if __name__ == "__main__":
    pass
